import cv2
import torch

from PIL import Image, ImageDraw

import dataset
import config as cfg


class Sailor:
    def __init__(self, path):
        self.size_in = cfg.input_size
        self.labels = dataset.classes()
        self.transform = dataset.transform()
        self.model = torch.load(path, map_location='cpu')
        print(self.model)

    def convert(self, image):
        print(image)  # [3,32,32]
        # 网络的输入为32X32
        # transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),torchvision.transforms.ToTensor()])
        image = self.transform(image)
        print(image.shape)
        # [3,32,32] --> [1,3,32,32]
        image = torch.reshape(image, [1, 3, self.size_in, self.size_in])
        return image

    def convert_cv(self, image):
        resized_img = cv2.resize(image, (self.size_in, self.size_in), interpolation=cv2.INTER_LINEAR)
        image_rgb = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)  # bgr-->rgb
        image = torch.from_numpy(image_rgb)  # 转torch.Tensor
        image = image.float()  # 输入的图片是整数，要和网络参数(浮点数)保持一致
        # print(image.shape)
        image = torch.reshape(image, [1, 3, self.size_in, self.size_in])  # [3,32,32] --> [1,3,32,32]  提升一张照片的维度
        return image

    def eval(self, img):
        self.model.eval()
        with torch.no_grad():
            output = self.model(img)
            idx = output.argmax(1)
            label = self.labels[idx]
        return label, output, idx

    def eval_file(self, img_path: str, show: bool):
        image_raw = Image.open(img_path)
        image = image_raw.convert('RGB')  # 将ARGB-->RGB
        image = self.convert(image)
        label, output, idx = self.eval(image)
        if show:
            draw = ImageDraw.Draw(image_raw)
            draw.text((10, 10), label, fill=(255, 0, 0))
            image_raw.show()
        print(output)
        print(idx.item())

    def eval_video(self, win_name: str):
        cap = cv2.VideoCapture(1)
        if not cap.isOpened():
            print("无法打开摄像头")
            exit()
        while True:
            ret, image_raw = cap.read()
            if not ret:
                print("无法获取帧")
                break
            image = self.convert_cv(image_raw)
            label, output, idx = self.eval(image)
            if idx >= 0:
                cv2.putText(image_raw, label, (10, 20), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255))
            cv2.imshow(win_name, image_raw)
            if cv2.waitKey(1) & 0xFF == ord('q'):  # 按下'q'键退出循环
                break
        cap.release()
        cv2.destroyAllWindows()
